import laspy
import open3d as o3d  
import numpy as np

def pc_normalize(pc):
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc


def replace_numbers(lst):  
    mapping = {16: 0, 27: 1, 28: 2}  
    return [mapping.get(num, 3) for num in lst] 



def repall(file):
    inFile = laspy.read(file)
    xyz = inFile.xyz  
    # 转换为numpy array  
    npy_xyz = np.array(xyz)  
    datan = pc_normalize(npy_xyz)
    new_cls = replace_numbers(inFile.classification)
    cls_array = np.array([cls for cls in zip(new_cls)])  
    reds =  inFile.red 
    greens = inFile.green
    blues =  inFile.blue

    # 首先，我们需要将16位整数转化为0-255的整数  
    reds = reds.astype(np.uint16) // 256  
    greens = greens.astype(np.uint16) // 256  
    blues = blues.astype(np.uint16) // 256  
    
    # 然后，我们可以将这三个通道的数据转化为RGB颜色  
    rgb = np.stack((reds, greens, blues), axis=-1)
    final_data = np.hstack([datan, rgb, cls_array])
    return final_data



import os  
  
# 定义文件夹路径  
folder_path = 'train_data/'  
  
# 使用os库检查文件夹是否存在  
if os.path.exists(folder_path):  
    # 使用os.listdir列出文件夹中的所有文件和子文件夹  
    for file_name in os.listdir(folder_path):  
        # 打印文件路径和名称  
        full_name= os.path.join(folder_path, file_name)
        print(full_name)
        array = repall(full_name)
        # indices_to_remove = np.where(array[:, 6] == 3)[0]  
  
        # # 计算需要删除的行数（即这些行的60%）  
        # num_rows_to_remove = int(len(indices_to_remove) * 0.2)  
        
        # # 根据计算出的行数，从原始数组中删除相应的行  
        # rows_to_keep = np.delete(indices_to_remove, np.arange(num_rows_to_remove))  
        # filtered_array = np.delete(array, rows_to_keep, axis=0)  
        np.save('proalltrain/%s.npy'%file_name[:-4], array)
else:  
    print("The folder does not exist.")



# array = repall("1/1_1_1.las")

# unique_vals, counts = np.unique(array[:,6], return_counts=True)  
  
# print("Unique values:", unique_vals)  
# print("Counts:", counts)

# 找到第七列中等于3的所有行的索引  
# indices_to_remove = np.where(array[:, 6] == 3)[0]  
  
# # 计算需要删除的行数（即这些行的60%）  
# num_rows_to_remove = int(len(indices_to_remove) * 0.2)  
  
# # 根据计算出的行数，从原始数组中删除相应的行  
# rows_to_keep = np.delete(indices_to_remove, np.arange(num_rows_to_remove))  
# filtered_array = np.delete(array, rows_to_keep, axis=0)  

# print(filtered_array.shape)

# unique_vals, counts = np.unique(filtered_array[:,6], return_counts=True)  
  
# print("Unique values:", unique_vals)  
# print("Counts:", counts)